# from gensim.test.utils import datapath, get_tmpfile
# from gensim.models import KeyedVectors
# from gensim.scripts.glove2word2vec import glove2word2vec
# import gensim
import os
import torch
from torch import nn
import dltools

#glove_file = datapath(r"D:\learn\深度学习\ai_learn\MNIST\glove.6B.100d.txt")
#word2vec_glove_file = get_tmpfile("glove.6B.100d.word2vec.temp")

#glove2word2vec(glove_file, word2vec_glove_file)
#model = KeyedVectors.load_word2vec_format(word2vec_glove_file)
#res = model.most_similar("banana")
#print(f"res:{res}")

#res = model.most_similar(positive=["man", "king"], negative=["woman"])
#print(f"res:{res}")

class TokenEmbedding:
    def __init__(self, file_path):
        self.idx_to_token, self.idx_to_vec = self._load_embedding(file_path)
        self.unknown_idx = 0
        self.token_to_idx = {token: idx for idx, token in enumerate(self.idx_to_token)}

    # 加载预训练词向量
    def _load_embedding(self, file_path):
        idx_to_token, idx_to_vec = ["<unk>"], []

        with open(os.path.join(file_path), "r", encoding="utf-8") as f:
            for line in f:
                elems = line.strip().split(" ")
                if len(elems) > 1:
                    token = elems[0]
                    vecs = [float(elem) for elem in elems[1:]]
                    idx_to_token.append(token)
                    idx_to_vec.append(vecs)
        idx_to_vec = [[0] * len(idx_to_vec[0])] + idx_to_vec
        # print(f"idx_to_vec:{idx_to_vec[:10]}")
        return idx_to_token, torch.tensor(idx_to_vec)

    def __getitem__(self, tokens):
        idxes = [self.token_to_idx.get(token, self.unknown_idx) for token in tokens]
        return self.idx_to_vec[torch.tensor(idxes)]

    def __len__(self):
        return len(self.idx_to_token)
"""
# 6b指的是语料库的大小是60亿
glove_6b100d = TokenEmbedding(r"./MNIST/glove.6B.100d.txt")
print(f"len(glove_6b100d): {len(glove_6b100d)}")

# 计算词的余弦相似度
def knn(W, x, k):
    print(f"x.shape:{x.shape}")
    cos = torch.mv(W, x.reshape(-1,)) / (torch.sqrt(torch.sum(W * W, axis=1) + 1e-9) * torch.sqrt((x*x).sum()))
    _, topk = torch.topk(cos, k=k)
    return topk, _

def get_similar_tokens(query_token, k, embed:TokenEmbedding ):
    topk, cos = knn(embed.idx_to_vec, embed[[query_token]], k+1)
    for i, c in zip(topk[1:], cos[1:]):
        print(f"{embed.idx_to_token[int(i)]}:cosine相似度:{float(c):.3f}")

get_similar_tokens("chip", 3, glove_6b100d)
print("----------")
get_similar_tokens("boy", 3, glove_6b100d)
print("----------")
get_similar_tokens("beautiful", 3, glove_6b100d)

def get_analogy(token_a, token_b, token_c, embed:TokenEmbedding):
    vecs = embed[[token_a, token_b, token_c]]
    x = vecs[1] - vecs[0] + vecs[2]
    topk, cos = knn(embed.idx_to_vec, x, 1)
    return embed.idx_to_token[topk[0]]

res = get_analogy("man", "woman", "son", glove_6b100d)
print(f"res:{res}")
res = get_analogy("queen", "king", "woman", glove_6b100d)
print(f"res:{res}")
"""